Knowledge Bases for Amazon Bedrock (with OpenSearch Serverless)をSAMで実装してみた
2024.05.26こんにちは、つくぼし(tsukuboshi0755)です！
Knowledge Bases for Amazon Bedrockが東京リージョンにやってきた事で、RAG構成の候補としてKnowledge Basesをより使いやすくなったかと思います。
以前Knowledge Basesを用いたRAG構成として、ベクターストアにAuroraを利用したパターンをCloudFormationで作成しました。
ただ多くの方々は、Knowledge Basesコンソールのクイック作成でも用いられる、ベクターストアにOpenSearch Serverlessを利用したパターンで作成する事が多いのではないでしょうか。
そこで今回は、ベクターストアにOpenSearch Serverlessを利用したパターンをSAMで実装してみたので、紹介したいと思います！
なおインデックスのセットアップの実施で一部カスタムリソースを含む形になるため、その点にご留意頂けるとありがたいです。
構成
今回SAMで作成する構成図は以下のとおりです。
Knowledge Bases for Amazon Bedrockに必要なリソースとして、以下のリソースも合わせて作成します。
- ベクターストア：OpenSearch Serverless
- データソース：S3
またLambda-Backed Custom Resourceを利用して、OpenSearch Serverlessに対するインデックスのセットアップを自動化します。
テンプレート
全体のコードは、以下のリポジトリに格納しています。
以下では、テンプレートの主要な部分を抜粋して説明します。
OpenSearch Serverless
OSSCollection:
Type: 'AWS::OpenSearchServerless::Collection'
Properties:
Name: !Sub '${AWS::StackName}-col'
Type: VECTORSEARCH
StandbyReplicas: !Ref OSSCollectionStandbyReplicas
DependsOn: OSSEncryptionPolicy
OSSEncryptionPolicy:
Type: 'AWS::OpenSearchServerless::SecurityPolicy'
Properties:
Name: !Sub '${AWS::StackName}-ep'
Type: encryption
Policy: !Sub >-
{
"Rules": [
{
"Resource": [
"collection/${AWS::StackName}-col"
],
"ResourceType": "collection"
}
],
"AWSOwnedKey": true
}
OSSNetworkPolicy:
Type: 'AWS::OpenSearchServerless::SecurityPolicy'
Properties:
Name: !Sub '${AWS::StackName}-np'
Type: network
Policy: !Sub >-
[
{
"Rules": [
{
"Resource": [
"collection/${AWS::StackName}-col"
],
"ResourceType": "dashboard"
},
{
"Resource": [
"collection/${AWS::StackName}-col"
],
"ResourceType": "collection"
}
],
"AllowFromPublic": true
}
]
OSSDataAccessPolicy:
Type: 'AWS::OpenSearchServerless::AccessPolicy'
Properties:
Name: !Sub '${AWS::StackName}-dp'
Type: data
Policy: !Sub
- >-
[
{
"Rules": [
{
"Resource": [
"collection/${AWS::StackName}-col"
],
"Permission": [
"aoss:CreateCollectionItems",
"aoss:UpdateCollectionItems",
"aoss:DescribeCollectionItems"
],
"ResourceType": "collection"
},
{
"Resource": [
"index/${AWS::StackName}-col/*"
],
"Permission": [
"aoss:CreateIndex",
"aoss:UpdateIndex",
"aoss:DescribeIndex",
"aoss:ReadDocument",
"aoss:WriteDocument"
],
"ResourceType": "index"
}
],
"Principal": [
"${BedrockKnowledgeBaseRoleArn}",
"${CreateIndexFunctionRoleArn}",
"${OSSDataAccessPrincipalArn}"
],
"Description": ""
}
]
- BedrockKnowledgeBaseRoleArn: !GetAtt BedrockKnowledgeBaseRole.Arn
CreateIndexFunctionRoleArn: !GetAtt CreateIndexFunctionRole.Arn
Knowledge Base for Amazon Bedrockのベクターストアとして必要なOpenSearch Serverlessを作成します。
今回のOpenSearch Serverlessは基本的に以下を参照して作成しているため、詳細についてはこちらのの記事をご参照ください。
なお上記のブログからの変更点として、データアクセスポリシーに対して後述のKnowledge Bases及びインデックスセットアップ関数に対するロールにも、アクセス権限を付与するようにしています。
S3 Bucket
DataSourceBucket:
Type: AWS::S3::Bucket
Properties:
BucketName: !Sub '${AWS::StackName}-ds-bucket-${AWS::AccountId}'
PublicAccessBlockConfiguration:
BlockPublicAcls: true
BlockPublicPolicy: true
IgnorePublicAcls: true
RestrictPublicBuckets: true
BucketEncryption:
ServerSideEncryptionConfiguration:
- ServerSideEncryptionByDefault:
SSEAlgorithm: AES256
BucketKeyEnabled: true
OwnershipControls:
Rules:
- ObjectOwnership: BucketOwnerEnforced
Knowledge Base for Amazon Bedrockのデータソースとして必要なS3バケットを作成します。
このバケットに対してユーザーがドキュメントをアップロードする事で、Knowledge Base for Amazon Bedrockがデータソースとして利用できるようになります。
IAM Policy
BedrockFMPolicyForKB:
Type: AWS::IAM::ManagedPolicy
Properties:
ManagedPolicyName: !Sub ${AWS::StackName}-bedrock-fm-policy
PolicyDocument:
Version: '2012-10-17'
Statement:
- Sid: BedrockInvokeModelStatement
Effect: Allow
Action:
- 'bedrock:InvokeModel'
Resource: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/${EmbeddingModelId}
Roles:
- !Ref BedrockKnowledgeBaseRole
BedrockOSSPolicyForKB:
Type: AWS::IAM::ManagedPolicy
Properties:
ManagedPolicyName: !Sub ${AWS::StackName}-bedrock-oss-policy
PolicyDocument:
Version: '2012-10-17'
Statement:
- Sid: OpenSearchServerlessAPIAccessAllStatement
Effect: Allow
Action:
- 'aoss:APIAccessAll'
Resource: !GetAtt OSSCollection.Arn
Roles:
- !Ref BedrockKnowledgeBaseRole
- !Ref CreateIndexFunctionRole
BedrockS3PolicyForKB:
Type: AWS::IAM::ManagedPolicy
Properties:
ManagedPolicyName: !Sub ${AWS::StackName}-bedrock-s3-policy
PolicyDocument:
Version: '2012-10-17'
Statement:
- Sid: S3ListBucketStatement
Effect: Allow
Action:
- 's3:ListBucket'
Resource: !GetAtt DataSourceBucket.Arn
Condition:
StringEquals:
aws:ResourceAccount: !Ref 'AWS::AccountId'
- Sid: S3GetObjectStatement
Effect: Allow
Action:
- 's3:GetObject'
Resource:
- !Join
- ''
- - 'arn:aws:s3:::'
- !Ref DataSourceBucket
- '/*'
Condition:
StringEquals:
aws:ResourceAccount: !Ref 'AWS::AccountId'
Roles:
- !Ref BedrockKnowledgeBaseRole
Lambda-Backed Custom Resource及びKnowledge Base for Amazon Bedrockに必要なIAMポリシーを作成します。
Bedrock、Secrets Manager、OpenSearch Serverless、S3に対するアクセス権限を設定し、IAMロールにアタッチします。
Lambda-Backed Custom Resource
- SAMテンプレート
SetupOSSIndex:
Type: 'Custom::OpenSearchServerlessIndex'
Properties:
ServiceToken: !GetAtt CreateIndexFunction.Arn
Region: !Ref 'AWS::Region'
Dimension: !If [IsTitanEmbedTextV1, 1536, 1024]
EmbeddingModelId: !Ref EmbeddingModelId
CollectionId: !GetAtt OSSCollection.Id
IndexName: !FindInMap [IndexMap, OSSIndexName, Name]
VectorField: !FindInMap [IndexMap, OSSVectorField, Name]
MappingText: !FindInMap [IndexMap, OSSMappingText, Name]
MappingMetadata: !FindInMap [IndexMap, OSSMappingMetadata, Name]
DependsOn:
- BedrockOSSPolicyForKB
CreateIndexFunctionRole:
Type: 'AWS::IAM::Role'
Properties:
RoleName: !Sub ${AWS::StackName}-create-index-function-role
AssumeRolePolicyDocument:
Version: '2012-10-17'
Statement:
- Effect: Allow
Principal:
Service: lambda.amazonaws.com
Action: 'sts:AssumeRole'
ManagedPolicyArns:
- arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
CreateIndexFunction:
Type: AWS::Serverless::Function
Properties:
FunctionName: !Sub '${AWS::StackName}-create-index-function'
Handler: index.lambda_handler
Role: !GetAtt CreateIndexFunctionRole.Arn
Runtime: python3.12
Timeout: 600
LoggingConfig:
LogFormat: JSON
ApplicationLogLevel: INFO
SystemLogLevel: INFO
Layers:
- !Ref OSSLayer
CodeUri: function/
OSSLayer:
Type: AWS::Serverless::LayerVersion
Properties:
LayerName: !Sub '${AWS::StackName}-create-index-function-layer'
ContentUri: layer/
CompatibleRuntimes:
- python3.12
Metadata:
BuildMethod: python3.12
- Lambdaコード
import logging
import time
from typing import Any, Dict
import boto3
import cfnresponse
from opensearchpy import AWSV4SignerAuth, OpenSearch, RequestsHttpConnection
logger = logging.getLogger()
def lambda_handler(event: Dict[str, Any], context: Any) -> None:
request_type = event["RequestType"]
logger.info(f"RequestType: {request_type}")
try:
if request_type == "Create":
# Get the parameters from the event
region = event["ResourceProperties"]["Region"]
dimension = event["ResourceProperties"]["Dimension"]
colection_id = event["ResourceProperties"]["CollectionId"]
index_name = event["ResourceProperties"]["IndexName"]
vector_field = event["ResourceProperties"]["VectorField"]
mapping_text = event["ResourceProperties"]["MappingText"]
mapping_metadata = event["ResourceProperties"]["MappingMetadata"]
# Create OpenSearch client
service = "aoss"
host = colection_id + "." + region + "." + service + "." + "amazonaws.com"
credentials = boto3.Session().get_credentials()
auth = AWSV4SignerAuth(credentials, region, service)
client = OpenSearch(
hosts=[{"host": host, "port": 443}],
http_auth=auth,
use_ssl=True,
verify_certs=True,
connection_class=RequestsHttpConnection,
pool_maxsize=20,
)
# Create OpenSearch Index (Wait for 30 seconds for index creation)
index_body = {
"settings": {
"index": {
"number_of_shards": "2",
"number_of_replicas": "0",
"knn": "true",
}
},
"mappings": {
"properties": {
vector_field: {
"type": "knn_vector",
"dimension": dimension,
"method": {
"name": "hnsw",
"space_type": "l2",
"engine": "faiss",
"parameters": {},
},
},
mapping_text: {
"type": "text",
},
mapping_metadata: {"type": "text", "index": "false"},
}
},
}
response = client.indices.create(
index_name,
body=index_body,
)
logger.info(f"Index creating: {response}")
time.sleep(30)
cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
if request_type == "Update":
cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
if request_type == "Delete":
cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
except Exception as e:
cfnresponse.send(event, context, cfnresponse.FAILED, {"Message": str(e)})
今回のテンプレートの肝となる箇所です。
Knowledge Base for Amazon BedrockのベクターストアとしてOpenSearch Serverlessを利用する際、別途インデックスのセットアップが必要となります。
こちらはCloudFormationでは対応していないため、カスタムリソースを使用してリクエストタイプがCreateとなる時のみ、OpenSearch Serverlessに対してOpenSearch APIをLambda関数で実行する事で、インデックスのセットアップを自動化します。
カスタムリソース及びリクエストタイプについては、以下の記事をご参照ください。
また今回はOpenSearch APIとしてopensearch-pyモジュールを使用しています。
そのためrequirements.txtにopensearch-pyを記載し、Lambda関数のデプロイ時にレイヤーとして追加しています。
opensearch-pyについては、以下の公式ドキュメントをご参照ください。
カスタムリソースとして定義したLambda関数では、以下のセットアップ処理についてOpenSearch APIを実行します。
- コレクションへの接続
- ベクターストア用のインデックス作成
カスタムリソースによりインデックスが正常に作成されているかどうかを確認するには、OpenSearchコンソールを使用します。
作成したコレクションに対して、以下のようなインデックスが確認できればOKです。
またコンソールだけでなく、awscurlまたはcurlコマンド等を用いる事で、インデックスの詳細設定を確認できます。
CloudShell等の環境で簡単に確認できるため、必要に応じてご参照ください。
# エンドポイントの指定
AOSS_ENDPOINT=<OpenSearch Serverlessのエンドポイント>
INDEX_NAME=bedrock-knowledge-base-default-index
# awscurlのインストール
pip install awscurl
# インデックスの確認
awscurl --service aoss --region ${AWS_REGION} \
-X GET \
${AOSS_ENDPOINT}/${INDEX_NAME} | jq -r .
上記を実施した際に、以下のようなインデックスの詳細設定が確認できればOKです。
{
"bedrock-knowledge-base-default-index": {
"aliases": {},
"mappings": {
"properties": {
"AMAZON_BEDROCK_METADATA": {
"type": "text",
"index": false
},
"AMAZON_BEDROCK_TEXT_CHUNK": {
"type": "text"
},
"bedrock-knowledge-base-default-vector": {
"type": "knn_vector",
"dimension": 1024, // ベクトルの次元数 1536 or 1024
"method": {
"engine": "faiss",
"space_type": "l2",
"name": "hnsw",
"parameters": {}
}
}
}
},
"settings": {
"index": {
"number_of_shards": "2",
"provided_name": "bedrock-knowledge-base-default-index",
"knn": "true",
"creation_date": "XXXXXXXXXXXXX",
"number_of_replicas": "0",
"uuid": "xxxxxxxxxxxxxxxxxxxx",
"version": {
"created": "XXXXXXXXX"
}
}
}
}
}
Knowledge Base for Amazon Bedrock
BedrockKnowledgeBaseRole:
Type: AWS::IAM::Role
Properties:
RoleName: !Sub ${AWS::StackName}-bedrock-kb-role
AssumeRolePolicyDocument:
Version: '2012-10-17'
Statement:
- Effect: Allow
Principal:
Service: [bedrock.amazonaws.com]
Action: ['sts:AssumeRole']
BedrockKnowledgeBase:
Type: AWS::Bedrock::KnowledgeBase
Properties:
Name: !Sub ${AWS::StackName}-knowledge-base
KnowledgeBaseConfiguration:
Type: VECTOR
VectorKnowledgeBaseConfiguration:
EmbeddingModelArn: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/${EmbeddingModelId}
RoleArn: !GetAtt BedrockKnowledgeBaseRole.Arn
StorageConfiguration:
Type: OPENSEARCH_SERVERLESS
OpensearchServerlessConfiguration:
CollectionArn: !GetAtt OSSCollection.Arn
FieldMapping:
MetadataField: !FindInMap [IndexMap, OSSMappingMetadata, Name]
TextField: !FindInMap [IndexMap, OSSMappingText, Name]
VectorField: !FindInMap [IndexMap, OSSVectorField, Name]
VectorIndexName: !FindInMap [IndexMap, OSSIndexName, Name]
DependsOn: SetupOSSIndex
BedrockKnowledgeBaseDS:
Type: AWS::Bedrock::DataSource
Properties:
DataDeletionPolicy: RETAIN
KnowledgeBaseId: !Ref BedrockKnowledgeBase
Name: !Sub ${AWS::StackName}-data-source
DataSourceConfiguration:
Type: S3
S3Configuration:
BucketArn: !GetAtt DataSourceBucket.Arn
最後にKnowledge Base for Amazon Bedrockを作成し、事前に作成したOpenSearch Serverless、Secrets Manager、S3を紐づけます。
EmbeddingModeArnでは、Knowledge Base for Amazon Bedrockに対して利用する埋め込みモデルのARNを指定します。
またOpensearchServerlessConfigurationでは、OpenSearch Serverlessに対するコレクション及びインデックス情報を指定します。
動作確認
上記のテンプレートをデプロイする事で、Knowledge Base for Amazon Bedrockを利用して正常に文書データを検索できるか確認します。
なお今回は、最近Knowledge Base for Amazon Bedrockが利用可能になった東京リージョンap-northeast-1でデプロイを行います。
SAM テンプレートのデプロイ
以下のコマンドを実行し、SAMテンプレートをデプロイします。
# プリンシパルARNを指定
IAM_PRINCIPAL_ARN=<IAMプリンシパルARN>
# SAM テンプレートのデプロイ
sam build
sam deploy --parameter-overrides OSSDataAccessPrincipalArn=${IAM_PRINCIPAL_ARN}
EmbeddingModelIdでは、Knowledge Base for Amazon Bedrockに対して利用する埋め込みモデルのIDを指定します。
今回はデフォルトのcohere.embed-multilingual-v3を指定します。
OSSCollectionStandbyReplicasは、OpenSearch Serverlessのスタンバイレプリカの有効/無効を指定します。
こちらも今回はデフォルトのDISABLEDを指定します。
OSSDataAccessPrincipalArnでは、OpenSearch Serverlessに対してアクセス権限を持つIAMプリンシパル(ユーザー/ロール)のARNを指定します。
こちらは指定必須であるため、OpenSearch Serverlessにアクセスが必要なIAMプリンシパルのARNを指定してください。
なおOpenSearch Serverlessのデプロイ時間が長く、およそ20分程度かかります。
S3バケットへのデータ投入
次にCloudShell上で以下のコマンドを実施し、PDFファイルを作成したS3バケットにアップロードします。
# S3 バケット名を設定
BUCKET_NAME=<バケット名>
# AWS の公式ドキュメントの PDF ファイルをダウンロード
mkdir -p pdf
cd pdf
wget https://docs.aws.amazon.com/ja_jp/amazondynamodb/latest/developerguide/dynamodb-dg.pdf -O DynamoDB.pdf
wget https://docs.aws.amazon.com/ja_jp/lambda/latest/dg/lambda-dg.pdf -O Lambda.pdf
wget https://docs.aws.amazon.com/ja_jp/vpc/latest/userguide/vpc-ug.pdf -O VPC.pdf
wget https://docs.aws.amazon.com/ja_jp/bedrock/latest/userguide/bedrock-ug.pdf -O Bedrock.pdf
cd ..
# S3 バケットに PDF ファイルをアップロード
aws s3 cp pdf s3://${BUCKET_NAME} --recursive
ファイルのアップロードが成功すると、以下の通り作成したS3バケットに4つのPDFファイルが作成されます。
Knowledge Baseデータソースの同期
続いてBedrockのコンソールに移動し、左ペインの"ナレッジベース"に切り替え、作成したナレッジベースをクリックします。
ナレッジベース画面中部あたりのデータソースに移動し、同期をクリックします。
同期中は、ステータスがSyncingになります。
データソースの同期にも、およそ20分程度かかります。
同期が完了すると、ステータスがReadyに戻り、最終同期時刻が表示されます。
Knowledge Baseでの検索実行
最後に検索が可能か確認するために、ナレッジベース画面のテストをクリックし、ナレッジベーステストを実施します。
例えば"回答を生成"をOFFにした状態で、「IPアドレス」と入力し送信すると、該当の用語に類似するデータソースのチャンクが表示されます。
また"回答を生成"をONにした状態で、「Lambdaに割り当てられる最大メモリを教えてください」と入力し送信すると、データソースのチャンクを元に正しく「10,240MB」という回答を返します。
これでナレッジベースが正しく動作している事を確認できました！
最後に
今回はベクターストアにOpenSearch Serverlessを利用したKnowledge Bases for Amazon BedrockをSAMで実装してみました。
ベクターストアのセットアップまで含めるとなかなかIaC化が難しいKnowledge Basesですが、コード管理する事でより検証がしやすくなるかと思います。
ぜひこちらのテンプレートを参考に、RAGの検証を試してみてください！
以上、つくぼし(tsukuboshi0755)でした！
![[アップデート] AWS Amplify AI kitが公開！ AI 機能を備えた Web アプリを迅速に構築できるようになりました](https://images.ctfassets.net/ct0aopd36mqt/wp-thumbnail-875a3d2bbdba43ad15684c9ee2e357b1/c77a61f7c646b6997c54cc55538ec251/aws-amplify)
